import pygame
import random

# 初始化pygame
pygame.init()

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GRAY = (128, 128, 128)

# 游戏屏幕大小
screen_width = 300
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Tetris")

# 设置时钟
clock = pygame.time.Clock()

# 方块尺寸
block_size = 30

# 定义方块形状
shapes = [
    [[1, 1, 1], [0, 1, 0]],  # I shape
    [[0, 2, 2], [2, 2, 0]],  # Z shape
    [[3, 3, 0], [0, 3, 3]],  # S shape
    [[4, 0, 0], [4, 4, 4]],  # T shape
    [[0, 0, 5], [5, 5, 5]],  # O shape
    [[6, 6, 6, 6]],  # L shape
    [[7, 7], [7, 7]]  # J shape
]

# 定义颜色
colors = [
    (0, 0, 0),
    (120, 37, 179),  # Purple
    (100, 179, 179),  # Light Blue
    (80, 34, 22),  # Brown
    (81, 51, 255),  # Dark Blue
    (255, 60, 122),  # Red
    (255, 138, 19),  # Orange
    (18, 118, 255)  # Light Green
]

# 游戏网格
grid = [[0 for x in range(10)] for y in range(20)]

# 当前方块的位置和形状
current_piece_x = 5
current_piece_y = 0
current_piece_shape = random.choice(shapes)
current_piece_color = random.choice(colors[1:])


# 下落函数
def move_down():
    global current_piece_y
    if not check_collision(grid, current_piece_shape, (current_piece_x, current_piece_y + 1)):
        current_piece_y += 1
    else:
        lock_piece_to_grid()


# 检查位置是否合法
def check_collision(board, shape, offset):
    off_x, off_y = offset
    for cy, row in enumerate(shape):
        for cx, cell in enumerate(row):
            if cell:
                if off_x + cx < 0 or off_x + cx >= len(board[0]):
                    return True
                if off_y + cy >= len(board):
                    return True
                if board[cy + off_y][cx + off_x]:
                    return True
    return False


# 锁定方块到游戏板上
def lock_piece_to_grid():
    global current_piece_shape, current_piece_color, current_piece_x, current_piece_y
    for i, row in enumerate(current_piece_shape):
        for j, cell in enumerate(row):
            if cell:
                grid[current_piece_y + i][current_piece_x + j] = current_piece_color
    clear_rows()
    generate_new_piece()


# 清除行
def clear_rows():
    global score
    num_rows_cleared = 0
    for i, row in enumerate(grid):
        if 0 not in row:
            num_rows_cleared += 1
            del grid[i]
            grid.insert(0, [0 for _ in range(10)])
            score += 10 * num_rows_cleared


# 生成新方块
def generate_new_piece():
    global current_piece_shape, current_piece_color, current_piece_x, current_piece_y
    current_piece_shape = random.choice(shapes)
    current_piece_color = random.choice(colors[1:])
    current_piece_x = 5
    current_piece_y = 0
    if check_collision(grid, current_piece_shape, (current_piece_x, current_piece_y)):
        game_over()


# 游戏结束
def game_over():
    global running
    running = False


# 移动方块
def move_left():
    global current_piece_x
    if not check_collision(grid, current_piece_shape, (current_piece_x - 1, current_piece_y)):
        current_piece_x -= 1


def move_right():
    global current_piece_x
    if not check_collision(grid, current_piece_shape, (current_piece_x + 1, current_piece_y)):
        current_piece_x += 1


# 旋转方块形状
def rotate_shape(shape):
    """
    顺时针旋转给定的方块形状90度。
    """
    # 转置矩阵
    transposed = list(zip(*shape))
    # 翻转每一行以实现顺时针旋转
    rotated = [list(row)[::-1] for row in transposed]
    return rotated


# 旋转当前方块
def rotate_current_piece():
    global current_piece_shape
    new_shape = rotate_shape(current_piece_shape)
    if not check_collision(grid, new_shape, (current_piece_x, current_piece_y)):
        current_piece_shape = new_shape


# 游戏主循环
running = True
score = 0

# 主循环
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_LEFT:
                move_left()
            elif event.key == pygame.K_RIGHT:
                move_right()
            elif event.key == pygame.K_DOWN:
                move_down()
            elif event.key == pygame.K_UP:
                rotate_current_piece()

    # 自动下降
    move_down()

    # 清屏
    screen.fill(BLACK)

    # 绘制网格上的方块
    for y in range(len(grid)):
        for x in range(len(grid[y])):
            if grid[y][x] != 0:
                pygame.draw.rect(screen, grid[y][x],
                                 [x * block_size, y * block_size, block_size, block_size])

    # 绘制当前方块
    for i in range(len(current_piece_shape)):
        for j in range(len(current_piece_shape[i])):
            if current_piece_shape[i][j] > 0:
                pygame.draw.rect(screen, current_piece_color,
                                 [(current_piece_x + j) * block_size,
                                  (current_piece_y + i) * block_size,
                                  block_size, block_size])

    # 更新屏幕
    pygame.display.flip()
    clock.tick(5)  # 控制游戏速度

pygame.quit()
